import torch
from torch.utils.data import Dataset,DataLoader

# 常规的做法，自己复写Dataset的get_item。然后用DataLoader进行处理
# torch 提供常用的Dataset

# Dataset可以进行预处理操作，
# 先读入到内存中
class MyStupidDataset(Dataset):
    def __init__(self):
        super(MyStupidDataset, self).__init__()
        self.data = torch.randn([1024, 10, 10])

    def __len__(self):
        return 1024

    def __getitem__(self, idx):
        return self.data[idx, :, :]


my_stupid_dataset=MyStupidDataset()

# - shuffle：设置为True的时候，每个世代都会打乱数据集
# - collate_fn：如何取样本的，我们可以定义自己的函数来准确地实现想要的功能
# - drop_last：告诉如何处理数据集长度除于batch_size余下的数据。True就抛弃，否则保留

# 普通的用法
# my_data_loader=DataLoader(my_stupid_dataset,batch_size=64,shuffle=True)
# for i,data in enumerate(my_data_loader):
#     print(i)

# 自定义回调，batch size中指定回调方法
def process(x):
    print(x)

my_data_loader_more=DataLoader(my_stupid_dataset,batch_size=64,shuffle=True,
                               collate_fn=process)

for i,data in enumerate(my_data_loader_more):
    print(i)